import numpy as np
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

def assign_cluster(x, c):
    distances = np.linalg.norm(x[:, np.newaxis] - c, axis=2)
    return np.argmin(distances, axis=1)

def Kmeans(data, k, epsilon=1e-4, iteration=100):
    n_samples, n_features = data.shape
    np.random.seed(0)
    centroids = data[np.random.choice(n_samples, k, replace=False)]

    for i in range(iteration):
        labels = assign_cluster(data, centroids)
        new_centroids = np.array([
            data[labels == j].mean(axis=0) if np.any(labels == j) else centroids[j]
            for j in range(k)
        ])
        shift = np.linalg.norm(new_centroids - centroids)
        if shift < epsilon:
            print(f"K-means算法在第 {i + 1} 次迭代后收敛 (质心移动距离: {shift:.6f})")
            break
        centroids = new_centroids

    return labels, centroids

if __name__ == "__main__":
    np.random.seed(42)
    data = np.vstack([
        np.random.randn(50, 2) + np.array([0, 0]),
        np.random.randn(50, 2) + np.array([5, 5]),
        np.random.randn(50, 2) + np.array([10, 0])
    ])

    Y, C = Kmeans(data, k=3)
    print("聚类结果标签：", Y)
    print("聚类中心：\n", C)

    plt.scatter(data[:, 0], data[:, 1], c=Y, cmap='viridis', s=30)
    plt.scatter(C[:, 0], C[:, 1], c='red', marker='X', s=200, label='Centroids')
    plt.legend()
    plt.title("K-Means 聚类结果")
    plt.show()
